import pyaudio
import wave
import datetime
import numpy as np
import matplotlib.pyplot as plt

# 音频输入参数
format = pyaudio.paInt16
channels = 1
rate = 44100
chunk_size = 1024
record_seconds = 2
cooldown_time = 1  # HIT 计数的冷却时间，单位为秒

# 创建一个PyAudio对象
audio = pyaudio.PyAudio()

# 打开音频输入流
stream = audio.open(format=format,
                    channels=channels,
                    rate=rate,
                    input=True,
                    frames_per_buffer=chunk_size)

# 上一次 HIT 的时间
last_hit_time = None

def calculate_threshold(audio_data):
    max_amplitude = np.max(audio_data)
    hit_threshold = int(0.5 * max_amplitude)

    print( f"threshold: {hit_threshold}" )
    return hit_threshold

sample_pool = []

def calculate_199_threshold(audioone,threshold):
    global sample_pool

    sample_pool.append(abs(audioone))

    if len(sample_pool) <= 199:
        return threshold
    else:
        reference_average = np.mean(sample_pool[-199:])
        dynamic_threshold = int(5.0 * reference_average)
        return dynamic_threshold
    

def count_hits(audio_data, threshold):
    hit_count = 0
    skip_count = 0

    for i, sample in enumerate(audio_data):


        dynamic_threshold = calculate_199_threshold(sample,threshold)

        # print(f"Dynamic Threshold: {i} {dynamic_threshold}")
        
        if skip_count > 0:
            skip_count -= 1
            continue

        if sample > threshold:
            print(f"HIT: {i}: {sample}")
            hit_count += 1
            skip_count = 199 # int(rate / chunk_size * cooldown_time)

        
        if sample > dynamic_threshold:
            print(f"Dyamic HIT: {i}: {sample}")
            hit_count += 1
            skip_count = 199 # int(rate / chunk_size * cooldown_time)

    return hit_count

def save_audio(filename, frames):
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(audio.get_sample_size(format))
        wf.setframerate(rate)
        wf.writeframes(b''.join(frames))

def plot_waveform(audio_data):
    # 创建时间轴
    time_axis = np.arange(0, len(audio_data)) / rate

    # 绘制波形图
    plt.figure()
    plt.plot(time_axis, audio_data)
    plt.xlabel("Time (s)")
    plt.ylabel("Amplitude")
    plt.title("Waveform")
    plt.show()

input("按回车键开始录制音频...")

print("录音开始...")
while True:
    frames = []
    for i in range(int(rate / chunk_size * record_seconds)):
        data = stream.read(chunk_size)
        frames.append(data)

    current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"audio_{current_time}.wav"
    save_audio(filename, frames)
    print(f"音频已保存为：{filename}")

    with wave.open(filename, "rb") as wf:
        num_frames = wf.getnframes()
        audio_data = np.frombuffer(wf.readframes(num_frames), dtype=np.int16)
    
    hit_threshold = calculate_threshold(audio_data)
    hit_count = count_hits(audio_data, hit_threshold)
    print(f"HIT次数：{hit_count}")
    
    plot_waveform(audio_data)

    user_input = input("按回车键继续录制音频，输入其他字符退出程序：")
    if user_input != "":
        break

    print("录音开始...")

print("录音结束...")
